import gymnasium as gym
import torch
import numpy as np
env = gym.make("CartPole-v1")
state, _ = env.reset()
print(type(state[0]))
print(state.shape)
state = np.expand_dims(state,0)
print(state.shape)
state = torch.tensor(state, dtype=torch.float)
print(state)
print(type(state))
print(state.shape)
next_state, reward, terminated, truncated, _ = env.step(0)
print(type(next_state[0]))
print(type(terminated))
print(type(reward))

q_values = torch.Tensor([[1.23, 4.56]])
action_index = q_values.argmax(dim=1)  # → tensor([1])
print(action_index)
print(type(action_index))
print(action_index.shape)
print(q_values.max(1)[0].shape)

import collections
import random

buffer = collections.deque(maxlen=5)
buffer.append((np.array([1,2,3]),2,3,4,5))
buffer.append((np.array([1,2,3]),2,3,4,5))
buffer.append((np.array([1,2,3]),2,3,4,5))
buffer.append((np.array([1,2,3]),2,3,4,5))
transitions = random.sample(buffer, 2)
states, actions, rewards, next_states, dones = zip(*transitions) # 用多重赋值接收迭代器
print(states)
print(type(states[0]))
actions = np.array(actions)
print(actions)
print(type(actions))
print(actions.shape)

x = torch.randn(4, 3)           # shape [4, 3]
idx = torch.tensor([0, 2])      # 默认是 int64 → OK
y = x.gather(1, idx.unsqueeze(1))  # ✅ 正确

# idx_int32 = torch.tensor([0, 2], dtype=torch.int32)
# y = x.gather(1, idx_int32.unsqueeze(1))  # ❌ RuntimeError!

